import torch
import numpy as np
from ultralytics import YOLO
from ultralytics.nn.modules.MSAD import ChannelAttention  # 确保该路径正确

def export_attn_weights(model_path, attn_weight_path, input_size=(3, 640, 640)):
    """
    载入预训练模型，使用一个示例输入对内部模型（model.model）做前向传播，
    遍历模型中所有 ChannelAttention 模块，提取它们在 forward 中保存的 attention_weights，
    并保存为 npz 文件。

    Args:
        model_path (str): 模型权重文件路径（例如 'best.pt'）。
        attn_weight_path (str): 要保存的 npz 文件路径（例如 'attn_weights.npz'）。
        input_size (tuple): 输入张量尺寸，默认 (3, 640, 640)。
    """
    # 加载YOLO模型（YOLO类包含训练流程，但我们只使用内部网络）
    model = YOLO(model_path)
    # 直接获取内部的网络部分，避免调用外层训练/预测接口
    underlying_model = model.model
    underlying_model.eval()  # 切换到推理模式

    # 获取模型所在设备，并将输入数据移到该设备上
    device = next(underlying_model.parameters()).device
    dummy_input = torch.randn(1, *input_size).to(device)

    # 使用纯推理前向传播，不调用 model.predict()（避免触发训练流程）
    with torch.no_grad():
        _ = underlying_model(dummy_input)

    # 遍历内部模型，提取所有 ChannelAttention 模块保存的 attention_weights
    attn_dict = {}
    for name, module in underlying_model.named_modules():
        if isinstance(module, ChannelAttention):
            if module.attention_weights is not None:
                attn_dict[name] = module.attention_weights

    if not attn_dict:
        print("未找到 ChannelAttention 的 attention_weights，请检查模块是否执行及保存正常。")
    else:
        np.savez(attn_weight_path, **attn_dict)
        print(f"Attention weights 已保存到 {attn_weight_path}")

if __name__ == '__main__':
    # 例如，使用预训练模型 best.pt 导出注意力权重文件 attn_weights.npz
    export_attn_weights("runs/train/exp/weights/best.pt", "attn_weights.npz")
